import cv2
import numpy as np
import matplotlib.pyplot as plt

def rgb_to_grayscale(image_path):
    """
    将彩色图像转换为灰度图像
    """
    # 读取彩色图像
    color_image = cv2.imread(image_path)
    if color_image is None:
        raise FileNotFoundError(f"无法找到图像文件: {image_path}")
    
    # 将BGR转换为灰度图
    gray_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2GRAY)
    
    # 确保图像大小为300x300
    if gray_image.shape != (300, 300):
        gray_image = cv2.resize(gray_image, (300, 300))
    
    return gray_image

def create_haar_filters(filter_size=7):
    """
    创建5个Haar滤波器
    filter_size: 滤波器大小，在5-10之间
    """
    filters = []
    
    # 1. 水平边缘特征滤波器 (黑白分界在中间)
    # 左边为正值，右边为负值
    filter1 = np.ones((filter_size, filter_size))
    mid = filter_size // 2
    filter1[:, mid+1:] = -1  # 右半部分设为-1
    filters.append(filter1)
    
    # 2. 垂直边缘特征滤波器 (上下分界)
    # 上半部分为正值，下半部分为负值
    filter2 = np.ones((filter_size, filter_size))
    filter2[mid+1:, :] = -1  # 下半部分设为-1
    filters.append(filter2)
    
    # 3. 水平线特征滤波器 (中间条状)
    # 中间行为正值，上下为负值
    filter3 = -np.ones((filter_size, filter_size))
    center_row = filter_size // 2
    filter3[center_row, :] = 1  # 中间行设为1
    filters.append(filter3)
    
    # 4. 垂直线特征滤波器 (中间条状)
    # 中间列为正值，左右为负值
    filter4 = -np.ones((filter_size, filter_size))
    center_col = filter_size // 2
    filter4[:, center_col] = 1  # 中间列设为1
    filters.append(filter4)
    
    # 5. 四矩形特征滤波器 (棋盘格模式)
    filter5 = np.ones((filter_size, filter_size))
    for i in range(filter_size):
        for j in range(filter_size):
            # 创建棋盘格模式：相邻区域符号相反
            if (i // (filter_size//2 + 1) + j // (filter_size//2 + 1)) % 2 == 0:
                filter5[i, j] = 1
            else:
                filter5[i, j] = -1
    filters.append(filter5)
    
    return filters

def manual_convolution(image, kernel):
    """
    手动实现卷积操作
    使用step-by-step的for循环方式
    """
    # 获取图像和滤波器尺寸
    img_h, img_w = image.shape
    kernel_h, kernel_w = kernel.shape
    
    # 计算填充大小以保持输出尺寸不变
    pad_h = kernel_h // 2
    pad_w = kernel_w // 2
    
    # 对图像进行零填充
    padded_image = np.pad(image, ((pad_h, pad_h), (pad_w, pad_w)), mode='constant')
    
    # 初始化输出特征图
    output = np.zeros((img_h, img_w))
    
    # step-by-step卷积计算
    for i in range(img_h):
        for j in range(img_w):
            # 提取当前卷积区域
            region = padded_image[i:i+kernel_h, j:j+kernel_w]
            # 执行逐元素相乘并求和
            output[i, j] = np.sum(region * kernel)
    
    return output

def normalize_to_uint8(feature_map):
    """
    将特征图归一化到0-255范围并转换为uint8类型
    """
    # 移除非有限值
    feature_map = np.nan_to_num(feature_map)
    
    # 归一化到0-255范围
    normalized = feature_map - np.min(feature_map)
    if np.max(normalized) > 0:
        normalized = normalized / np.max(normalized) * 255
    
    return normalized.astype(np.uint8)

def visualize_results(original, feature_maps, filter_names):
    """
    可视化原始图像和特征图
    """
    plt.figure(figsize=(15, 10))
    
    # 显示原始图像
    plt.subplot(2, 3, 1)
    plt.imshow(original, cmap='gray')
    plt.title('Original Grayscale Image')
    plt.axis('off')
    
    # 显示5个特征图
    for i in range(5):
        plt.subplot(2, 3, i+2)
        plt.imshow(feature_maps[i], cmap='gray')
        plt.title(f'Feature Map {i+1}: {filter_names[i]}')
        plt.axis('off')
    
    plt.tight_layout()
    plt.show()

def main():
    # Image path
    image_path = 'jay.png'
    
    try:
        # 1. 将彩色图像转换为灰度图像
        print("Step 1: Converting color image to grayscale...")
        gray_image = rgb_to_grayscale(image_path)
        print(f"Grayscale image dimensions: {gray_image.shape}")
        
        # 2. 创建 Harr 滤波器
        print("Step 2: Creating Haar filters...")
        filter_size = 7  # Using 7x7 filter size
        haar_filters = create_haar_filters(filter_size)
        filter_names = [
            "Horizontal Edge",
            "Vertical Edge", 
            "Horizontal Line",
            "Vertical Line",
            "Four Rectangle"
        ]
        
        # 3. 手动计算卷积
        print("Step 3: Performing manual convolution...")
        feature_maps = []
        for i, kernel in enumerate(haar_filters):
            print(f"  Processing filter {i+1}: {filter_names[i]}")
            feature_map = manual_convolution(gray_image, kernel)
            feature_maps.append(feature_map)
        
        # 4. 标准化特征图
        print("Step 4: Normalizing feature maps...")
        normalized_maps = []
        for i, feature_map in enumerate(feature_maps):
            normalized_map = normalize_to_uint8(feature_map)
            normalized_maps.append(normalized_map)
        
        # 5. 保存特征图
        print("Step 5: Saving feature maps...")
        for i, normalized_map in enumerate(normalized_maps):
            filename = f'haar_feature_{i+1}_{filter_names[i].replace(" ", "_").lower()}.png'
            cv2.imwrite(filename, normalized_map)
            print(f"  Saved: {filename}")
        
        # 6. 可视化结果
        print("Step 6: Visualizing results...")
        visualize_results(gray_image, normalized_maps, filter_names)
        
        print("All steps completed successfully!")
        
    except Exception as e:
        print(f"Error occurred: {e}")

if __name__ == "__main__":
    main()